# -*- coding: utf-8 -*-
from tensorflow.keras.callbacks import Callback
import numpy as np
from skimage.util import img_as_ubyte
from skimage import io as skio
import os


class MyCallback(Callback):
    """Keras (version=2.3.1) Callback 编写模板"""

    def __init__(self, model=None, data=None, label=None, metric=None, save_path=None):
        super().__init__()
        self.data = data
        self.label = label
        self.metric = metric
        self.model = model
        self.save_path = save_path

    def on_train_begin(self, logs: dict):
        """在整个训练开始时会调用次函数

            Parameters:
            ----------
                logs: dict, 该参数在当前版本默认为 None, 主要是为未来的 keras 版本的新行为预留位置
        """
        pass

    def on_train_end(self, logs: dict):
        """在整个训练结束时调用次函数

            Parameters:
            ----------
                logs: dict, 该参数在当前版本默认为 None, 主要是为未来的 keras 版本的新行为预留位置
        """
        pass

    def on_epoch_begin(self, epoch, logs: dict):
        """在每个 epoch 开始的时候调用此函数

            Parameters:
            ----------
                epoch: int, 当前为第几个 epoch, 从 1 开始
                logs: dict, 为空
        """
        #print('On epoch begin', epoch, logs)
        pass

    def on_epoch_end(self, epoch, logs: dict):
        """在每个 epoch 结束的时候调用此函数

            Parameters:
            ----------
                epoch: int, 当前为第几个 epoch, 从 1 开始
                logs: dict, 包含了当前 epoch 的一些信息，主要的 key 有:
                    - accuracy
                    - loss
                    - val-accuracy（仅在 fit 中开启 validation 时才有）
                    - val-loss（仅在 fit 中开启 validation 时才有）
        """
        print('epoch: {} end, {}'.format(epoch, logs))
        preds = self.model.predict(self.data)
        img = self.combine_img_prediction(self.data, self.label, preds)
        img = img_as_ubyte(img)
        mean_major_iou = 0
        mean_sec_iou = 0

        major_mean_precision = 0
        major_mean_recall = 0

        major_mean_f1 = 0

        sec_mean_precision = 0
        sec_mean_recall = 0

        sec_mean_f1 = 0

        for index, item in enumerate(self.label):
            res = self.metric((self.label[index] > 0).astype(np.uint8), preds[index])
            mean_major_iou += res['major_iou']
            mean_sec_iou += res['sec_iou']
            major_mean_precision += res['report']['0']['precision']
            major_mean_recall += res['report']['0']['recall']
            major_mean_f1 += res['report']['0']['f1-score']

            sec_mean_precision += res['report']['1']['precision']
            sec_mean_recall += res['report']['1']['recall']
            sec_mean_f1 += res['report']['1']['f1-score']

        mean_major_iou /= self.label.shape[0]
        mean_sec_iou /= self.label.shape[0]
        major_mean_precision /= self.label.shape[0]
        major_mean_recall /= self.label.shape[0]

        major_mean_f1 /= self.label.shape[0]

        sec_mean_precision /= self.label.shape[0]
        sec_mean_recall /= self.label.shape[0]

        sec_mean_f1 /= self.label.shape[0]

        mean_major_iou = mean_major_iou.numpy()
        mean_sec_iou = mean_sec_iou.numpy()
        print("major vein val data iou : {}, sec vein mean iou: {}".format(mean_major_iou, mean_sec_iou))
        #skio.imsave("/data/seg/prediction-unet_resnet50/epoch_{}.png".format(epoch-1), img)
        skio.imsave(os.path.join(self.save_path, "epoch_{}.png".format(epoch)), img)

    def on_batch_begin(self, batch, logs: dict):
        """在每个 batch 开始的时候调用此函数

            Parameters:
            ----------
                batch: int, 当前为第几个 batch, 从 1 开始
                logs: dict, 包含了当前 batch 的一些信息，主要的 key 有:
                    - batch: 同参数 batch
                    - size: batch 的大小
        """
        pass

    def on_batch_end(self, batch, logs: dict):
        """在每个 batch 结束的时候调用此函数

            Parameters:
            ----------
                batch: int, 当前为第几个 batch, 从 1 开始
                logs: dict, 包含了当前 batch 的一些信息，主要的 key 有:
                    - batch: 同参数 batch
                    - size: batch 的大小
                    - loss
                    - accuracy（仅当启用了 acc 监视）
        """
        # print(batch, logs)

    def combine_img_prediction(self, data, gt, pred):
        """
        Combines the data, grouth thruth and the prediction into one rgb image

        :param data: the data tensor
        :param gt: the ground thruth tensor
        :param pred: the prediction tensor

        :returns img: the concatenated rgb image
        """
        data = data[0:5]
        gt = gt[0:5]
        pred = pred[0:5]
        ny = int(pred.shape[2])
        ch = int(gt.shape[3])
        img = np.zeros([data.shape[0] * ny, ny * 3, 3])
        for i in range(data.shape[0]):
            img[i * ny:(i + 1) * ny, 0:ny, 0:3] = data[i]
            img[i * ny:(i + 1) * ny, ny:2 * ny, 0:3] = gt[i]
            img[i * ny:(i + 1) * ny, 2 * ny:3 * ny, 0:3] = pred[i]
        return img[:, :, 0:3]
